package com.corn.flink.lesson6;

import cn.hutool.core.collection.CollUtil;
import com.corn.flink.lesson4.Event;
import lombok.Data;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

/**
 * @author JimWu
 * @date 2023/3/24 10:44
 **/
public class TopNDemo {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
//        env.setParallelism(1);
        List<String> eventType = CollUtil.newArrayList("add", "delete", "update", "select");
        DataStreamSource<Event> ds = env.addSource(new RichSourceFunction<Event>() {
            private boolean flag = true;

            @Override
            public void run(SourceContext<Event> ctx) throws Exception {
                while (flag) {
                    ctx.collect(new Event(
                            eventType.get(ThreadLocalRandom.current().nextInt(eventType.size())),
                            System.currentTimeMillis()
                    ));
                    TimeUnit.MILLISECONDS.sleep(500);
                }
            }

            @Override
            public void cancel() {
                this.flag = false;
            }
        });
        ds.assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                    @Override
                    public long extractTimestamp(Event element, long recordTimestamp) {
                        return element.timestamp;
                    }
                }))
                .keyBy(data -> data.type)
                .window(TumblingEventTimeWindows.of(Time.seconds(5)))
                .aggregate(new AggregateFunction<Event, Long, Long>() {
                    @Override
                    public Long createAccumulator() {
                        return 0L;
                    }

                    @Override
                    public Long add(Event value, Long accumulator) {
                        return accumulator + 1;
                    }

                    @Override
                    public Long getResult(Long accumulator) {
                        return accumulator;
                    }

                    @Override
                    public Long merge(Long a, Long b) {
                        return null;
                    }
                }, new ProcessWindowFunction<Long, ViewCount, String, TimeWindow>() {
                    @Override
                    public void process(String s, Context context, Iterable<Long> elements, Collector<ViewCount> out) throws Exception {
                        Long amount = elements.iterator().next();
//                        Timestamp startTime = new Timestamp(context.window().getStart());
//                        Timestamp endTime = new Timestamp(context.window().getEnd());
                        ViewCount viewCount = new ViewCount();
                        viewCount.setCount(amount);
                        viewCount.setType(s);
                        viewCount.setEndTimestamp(context.window().getEnd());
                        out.collect(viewCount);
                    }
                })
                // 计算按窗口结束时间进行分区 统计每个窗口时间段下的topN
                .keyBy(data -> data.endTimestamp)
                .process(new CountTopKeyProcessFunction(2))
                .print();

        env.execute();
    }

    static class CountTopKeyProcessFunction extends KeyedProcessFunction<Long, ViewCount, String> {

        private final Integer top;

        private ListState<ViewCount> listState;

        private final String listStateName = "type-count-list-state";

        CountTopKeyProcessFunction(Integer top) {
            this.top = top;
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            // 从上下文获取状态句柄
            listState = getRuntimeContext().getListState(new ListStateDescriptor<ViewCount>(listStateName, Types.POJO(ViewCount.class)));
        }

        @Override
        public void processElement(ViewCount value, Context ctx, Collector<String> out) throws Exception {
            // 将数据保存到状态列表中
            this.listState.add(value);
            // 注册一个定时器 window end + 1 触发onTime方法计算出TOP N 的数据
            ctx.timerService().registerEventTimeTimer(value.endTimestamp + 5);
        }

        @Override
        public void onTimer(long timestamp, OnTimerContext ctx, Collector<String> out) throws Exception {
            // 获取到当前状态句柄的数据进行排序
//            this.listState.get(
            List<ViewCount> topNList = StreamSupport.stream(this.listState.get().spliterator(), false)
                    .sorted(Comparator.comparingLong(ViewCount::getCount).reversed())
                    .limit(top)
                    .collect(Collectors.toList());
            out.collect(topNList.toString());
            this.listState.clear();
        }
    }

    @Data
    public static class ViewCount {
        public String type;

        public Long count;

        public Long endTimestamp;
    }
}
